#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
设备编码生成器

根据业务规则生成各类设备的唯一编码：
- 台区：T + 9位省市区码（GB/T 2260 取到区县）+ 5位日级流水，长度15
- 电箱：B + 台区码 + 4位流水，长度19
- RTU：R + 电箱码 + 3位流水，长度23（或 R + 台区码 + 5位流水）
- 电表：M + RTU码 + 4位流水，长度28

特点：
- Redis预取号段+本地缓冲机制
- 支持灾难恢复（MySQL持久化）
- 并发安全+容量预警
"""

import asyncio
import threading
from datetime import date
from queue import Queue, Empty
from typing import Dict, Any

from common.database import DatabaseService
from common.logger import create_logger
from common.redis import RedisService, get_redis_service
from app.common.code_persistence import CodeSerialPersistence

logger = create_logger(__name__)


class CodeGenerationError(Exception):
	"""编码生成异常"""
	pass


class DeviceCodeGenerator:
	"""设备编码生成器"""

	# 编码规则常量
	STATION_PREFIX = "T"
	BOX_PREFIX = "B"
	RTU_PREFIX = "R"
	METER_PREFIX = "M"

	STATION_CODE_LEN = 15  # T + 9位区码 + 5位流水
	BOX_CODE_LEN = 19  # B + 台区码(14位) + 4位流水
	RTU_CODE_LEN = 16  # R + 9位区码 + 3位台区流水 + 3位RTU流水
	METER_CODE_LEN = 20  # M + 9位区码 + 3位台区流水 + 3位RTU流水 + 4位电表流水

	# 流水号池配置
	DAILY_POOL_SIZE = 100000  # 台区日级流水号池大小
	BOX_POOL_SIZE = 9999  # 电箱流水号池大小
	RTU_POOL_SIZE = 9999  # RTU流水号池如小（修改为9999）
	METER_POOL_SIZE = 9999  # 电表流水号池大小

	# 本地缓存配置
	LOCAL_CACHE_SIZE = 100  # 本地缓存号段大小（修改为100避免与RTU池大小冲突）
	REFILL_THRESHOLD = 0.2  # 剩余20%时预取

	def __init__(self):
		self.redis_service: RedisService = get_redis_service()
		self._local_caches: Dict[str, Queue] = {}
		self._cache_locks: Dict[str, threading.Lock] = {}

	def _get_local_cache(self, pool_key: str) -> Queue:
		"""获取指定池的本地缓存队列"""
		if pool_key not in self._local_caches:
			self._local_caches[pool_key] = Queue()
			self._cache_locks[pool_key] = threading.Lock()
		return self._local_caches[pool_key]

	def _get_cache_lock(self, pool_key: str) -> threading.Lock:
		"""获取指定池的缓存锁"""
		if pool_key not in self._cache_locks:
			self._cache_locks[pool_key] = threading.Lock()
		return self._cache_locks[pool_key]

	async def _get_next_sequence(self, pool_key: str, max_value: int, db: DatabaseService = None) -> int:
		"""
		获取下一个序列号

		Args:
			pool_key: Redis池键名
			max_value: 最大值限制
			db: 数据库服务实例，用于降级

		Returns:
			序列号

		Raises:
			CodeGenerationError: 编码生成失败
		"""
		cache = self._get_local_cache(pool_key)
		lock = self._get_cache_lock(pool_key)

		with lock:
			# 先尝试从本地缓存获取
			try:
				if not cache.empty():
					sequence = cache.get_nowait()

					# 检查是否需要预取
					if cache.qsize() < self.LOCAL_CACHE_SIZE * self.REFILL_THRESHOLD:
						# 异步预取，不阻塞当前请求
						asyncio.create_task(self._refill_local_cache(pool_key, max_value))

					return sequence
			except Empty:
				pass

			# 本地缓存为空，直接从Redis批量获取
			sequences = await self._fetch_sequences_from_redis(pool_key, max_value, self.LOCAL_CACHE_SIZE, db)

			if not sequences:
				raise CodeGenerationError(f"无法获取序列号: {pool_key}")

			# 将序列号放入本地缓存
			first_seq = sequences[0]
			for seq in sequences[1:]:
				cache.put(seq)

			return first_seq

	async def _refill_local_cache(self, pool_key: str, max_value: int, db: DatabaseService = None):
		"""异步补充本地缓存"""
		try:
			sequences = await self._fetch_sequences_from_redis(pool_key, max_value, self.LOCAL_CACHE_SIZE, db)
			if sequences:
				cache = self._get_local_cache(pool_key)
				lock = self._get_cache_lock(pool_key)
				with lock:
					for seq in sequences:
						if not cache.full():
							cache.put(seq)
				logger.debug(f"补充本地缓存成功: {pool_key}, 数量: {len(sequences)}")
		except Exception as e:
			logger.warning(f"补充本地缓存失败: {pool_key}, error: {e}")

	async def _fetch_sequences_from_redis(self, pool_key: str, max_value: int, count: int, db: DatabaseService = None) -> list[int]:
		"""
		从Redis批量获取序列号

		Args:
			pool_key: Redis池键名
			max_value: 最大值限制
			count: 获取数量
			db: 数据库服务实例，用于降级

		Returns:
			序列号列表
		"""
		if not self.redis_service.is_available():
			# Redis不可用，降级到数据库
			return await self._fetch_sequences_from_db(pool_key, max_value, count, db)

		# 使用Lua脚本保证原子性：INCR + 上限检查 + 告警
		lua_script = """
        local key = KEYS[1]
        local max_val = tonumber(ARGV[1])
        local count = tonumber(ARGV[2])
        local warning_threshold = tonumber(ARGV[3])

        local current = redis.call('GET', key)
        if not current then
            current = 0
        else
            current = tonumber(current)
        end

        -- 检查是否超过上限
        if current + count > max_val then
            return {-1, current, max_val}  -- 返回错误标识
        end

        -- 批量递增
        local new_val = redis.call('INCRBY', key, count)

        -- 检查是否达到预警阈值
        local warning_flag = 0
        if new_val >= max_val * warning_threshold then
            warning_flag = 1
        end

        -- 返回起始值和结束值
        local start_val = new_val - count + 1
        return {start_val, new_val, warning_flag}
        """

		try:
			if not self.redis_service.is_available():
				return await self._fetch_sequences_from_db(pool_key, max_value, count, db)

			warning_threshold = 0.9  # 90%预警阈值
			result = self.redis_service.eval_lua_script(lua_script, [pool_key], [max_value, count, warning_threshold])

			if result[0] == -1:
				# 达到上限，使用Lua脚本返回的current值作为Redis当前值
				current, max_val = result[1], result[2]
				logger.warning(f"编码池 {pool_key} 已达上限: {current}/{max_val}，降级到数据库")
				# 降级到数据库，使用Lua脚本返回的current值
				return await self._fetch_sequences_from_db(pool_key, max_value, count, db, current)

			start_val, end_val, warning_flag = result

			# 检查预警
			if warning_flag == 1:
				logger.warning(f"编码池 {pool_key} 使用量已达90%: {end_val}/{max_value}")
				# TODO: 发送告警通知

			# 生成序列号列表
			sequences = list(range(start_val, end_val + 1))
			return sequences

		except Exception as e:
			logger.error(f"从Redis获取序列号失败: {pool_key}, error: {e}")
			# 降级到数据库，尝试获取Redis当前值作为起始值
			redis_current_value = self._get_redis_current_value(pool_key)
			return await self._fetch_sequences_from_db(pool_key, max_value, count, db, redis_current_value)

	def _get_redis_current_value(self, pool_key: str) -> int:
		"""
		获取Redis中的当前值

		Args:
			pool_key: 池键名

		Returns:
			当前值，如果获取失败返回0
		"""
		try:
			if not self.redis_service.is_available():
				return 0

			value = self.redis_service.client.get(pool_key)
			if value is None:
				return 0

			return int(value)
		except Exception as e:
			logger.warning(f"获取Redis当前值失败: {pool_key}, error: {e}")
			return 0

	async def _fetch_sequences_from_db(self, pool_key: str, max_value: int, count: int, db: DatabaseService = None, redis_current_value: int = None) -> list[int]:
		"""
		从数据库获取序列号（Redis不可用时的降级方案）

		Args:
			pool_key: 池键名
			max_value: 最大值限制
			count: 获取数量
			db: 数据库服务实例
			redis_current_value: Redis中的当前值，用于同步数据库初始值

		Returns:
			序列号列表
		"""

		try:
			# 获取数据库连接
			if db is None:
				raise CodeGenerationError("数据库降级需要外部传入DatabaseService实例")

			# 解析池键获取池信息
			pool_info = self._parse_pool_key(pool_key)

			# 检查是否已在事务中
			in_transaction = not getattr(db.conn, 'autocommit', True)

			if in_transaction:
				# 已在事务中,直接执行不开启新事务
				# 获取当前偏移量
				offset_info = await CodeSerialPersistence.get_pool_offset(db, pool_key)

				if not offset_info:
					# 首次创建池，如果Redis有当前值，使用Redis的值作为起始值
					current_max = redis_current_value if redis_current_value is not None else 0
					if redis_current_value is not None:
						logger.info(f"数据库池首次创建，同步Redis当前值: {pool_key}, redis_value={redis_current_value}")
					await CodeSerialPersistence.upsert_pool_offset(
						db=db,
						pool_name=pool_key,
						pool_type=pool_info["type"],
						max_value=current_max + count,
						pool_size=max_value,
						area_code=pool_info.get("area_code"),
						parent_code=pool_info.get("parent_code"),
						date_str=pool_info.get("date_str")
					)
				else:
					current_max = offset_info["max_value"]

					# 检查是否超过限制
					if current_max + count > max_value:
						raise CodeGenerationError(f"数据库池 {pool_key} 已达上限: {current_max}/{max_value}")

					# 更新偏移量
					await CodeSerialPersistence.upsert_pool_offset(
						db=db,
						pool_name=pool_key,
						pool_type=pool_info["type"],
						max_value=current_max + count,
						pool_size=max_value,
						area_code=pool_info.get("area_code"),
						parent_code=pool_info.get("parent_code"),
						date_str=pool_info.get("date_str")
					)

				# 生成序列号列表
				start_val = current_max + 1
				end_val = current_max + count
				sequences = list(range(start_val, end_val + 1))

				logger.info(f"数据库降级获取序列号成功(已在事务中): {pool_key}, 范围: {start_val}-{end_val}")
				return sequences
			else:
				# 未在事务中,开启新事务
				async with db.transaction():
					# 获取当前偏移量
					offset_info = await CodeSerialPersistence.get_pool_offset(db, pool_key)

					if not offset_info:
						# 首次创建池，如果Redis有当前值，使用Redis的值作为起始值
						current_max = redis_current_value if redis_current_value is not None else 0
						if redis_current_value is not None:
							logger.info(f"数据库池首次创建，同步Redis当前值: {pool_key}, redis_value={redis_current_value}")
						await CodeSerialPersistence.upsert_pool_offset(
							db=db,
							pool_name=pool_key,
							pool_type=pool_info["type"],
							max_value=current_max + count,
							pool_size=max_value,
							area_code=pool_info.get("area_code"),
							parent_code=pool_info.get("parent_code"),
							date_str=pool_info.get("date_str")
						)
					else:
						current_max = offset_info["max_value"]

						# 检查是否超过限制
						if current_max + count > max_value:
							raise CodeGenerationError(f"数据库池 {pool_key} 已达上限: {current_max}/{max_value}")

						# 更新偏移量
						await CodeSerialPersistence.upsert_pool_offset(
							db=db,
							pool_name=pool_key,
							pool_type=pool_info["type"],
							max_value=current_max + count,
							pool_size=max_value,
							area_code=pool_info.get("area_code"),
							parent_code=pool_info.get("parent_code"),
							date_str=pool_info.get("date_str")
						)

					# 生成序列号列表
					start_val = current_max + 1
					end_val = current_max + count
					sequences = list(range(start_val, end_val + 1))

					logger.info(f"数据库降级获取序列号成功: {pool_key}, 范围: {start_val}-{end_val}")
					return sequences

		except Exception as e:
			logger.error(f"数据库降级获取序列号失败: {pool_key}, error: {e}")
			raise CodeGenerationError(f"数据库降级失败: {e}")

	@staticmethod
	def _validate_area_code(area_code: str) -> bool:
		"""
		验证区县码是否符合GB/T 2260标准

		Args:
			area_code: 9位区县码

		Returns:
			是否有效
		"""
		if not area_code or len(area_code) != 9:
			return False

		if not area_code.isdigit():
			return False

		# 基本格式检查：省市区都不能为00
		province = area_code[:2]
		city = area_code[2:4]
		area = area_code[4:6]

		if province == "00" or city == "00" or area == "00":
			return False

		# TODO: 可以添加更严格的GB/T 2260标准验证
		return True

	def _parse_pool_key(self, pool_key: str) -> Dict[str, Any]:
		"""
		解析池键获取类型和参数

		Args:
			pool_key: 池键名

		Returns:
			解析结果字典
		"""
		parts = pool_key.split(":")
		result = {"type": parts[0].upper()}

		if len(parts) >= 2:
			if parts[0] == "station":
				# station:{area_code}:{date_str}
				result["area_code"] = parts[1]
				if len(parts) >= 3:
					result["date_str"] = parts[2]
			elif parts[0] in ["box", "rtu", "meter"]:
				# box:{station_code}, rtu:{parent_code}, meter:{rtu_code}
				result["parent_code"] = parts[1]
				if len(parts) >= 3:
					result["parent_type"] = parts[2]

		return result

	async def _get_area_code_from_address(self, db: DatabaseService, address_id: int) -> str:
		"""
		通过address_id获取9位区县码

		Args:
			db: 数据库服务
			address_id: 地址ID

		Returns:
			9位区县码

		Raises:
			CodeGenerationError: 获取失败
		"""
		try:
			# 查询地址信息
			sql = """
                SELECT region_code, level, name, full_name
                FROM base_region_address
                WHERE id = %s AND is_active = 1
            """
			address = db.get_one(sql, (address_id,))

			if not address:
				raise CodeGenerationError(f"地址ID {address_id} 不存在或已停用")

			region_code = address.get("region_code")
			if not region_code:
				raise CodeGenerationError(f"地址ID {address_id} 缺少区域码")

			# 确保是9位区县码格式
			if len(region_code) == 6:
				# 6位标准码，补充3位到9位：前6位 + 000
				area_code = region_code + "000"
			elif len(region_code) >= 9:
				# 取前9位作为区县码
				area_code = region_code[:9]
			else:
				raise CodeGenerationError(f"地址ID {address_id} 区域码格式无效: {region_code}")

			if not self._validate_area_code(area_code):
				raise CodeGenerationError(f"区县码格式无效: {area_code}")

			return area_code

		except Exception as e:
			if isinstance(e, CodeGenerationError):
				raise
			logger.exception(f"获取区县码失败: address_id={address_id}")
			raise CodeGenerationError(f"获取区县码失败: {e}")

	# ============ 台区编码生成 ============

	async def generate_station_code(self, db: DatabaseService, address_id: int) -> str:
		"""
		生成台区编码: T + 9位区县码 + 5位日级流水

		Args:
			db: 数据库服务
			address_id: 地址ID

		Returns:
			台区编码

		Raises:
			CodeGenerationError: 生成失败
		"""
		try:
			# 1. 获取区县码
			area_code = await self._get_area_code_from_address(db, address_id)

			# 2. 生成池键：station:{area_code}:{yyyyMMdd}
			today = date.today().strftime("%Y%m%d")
			pool_key = f"station:{area_code}:{today}"

			# 3. 获取流水号
			sequence = await self._get_next_sequence(pool_key, self.DAILY_POOL_SIZE, db)

			# 4. 组装编码
			station_code = f"{self.STATION_PREFIX}{area_code}{sequence:05d}"

			logger.info(f"生成台区编码: {station_code}, area_code={area_code}, sequence={sequence}")
			return station_code

		except Exception as e:
			if isinstance(e, CodeGenerationError):
				raise
			logger.exception(f"生成台区编码失败: address_id={address_id}")
			raise CodeGenerationError(f"生成台区编码失败: {e}")

	# ============ 电箱编码生成 ============

	async def generate_box_code(self, station_code: str, db: DatabaseService = None) -> str:
		"""
		生成电箱编码: B + 台区码(去掉T前缀) + 4位流水

		Args:
			station_code: 台区编码
			db: 数据库服务实例，用于数据库降级

		Returns:
			电箱编码

		Raises:
			CodeGenerationError: 生成失败
		"""
		try:
			if not station_code or len(station_code) != self.STATION_CODE_LEN or not station_code.startswith(
					self.STATION_PREFIX):
				raise CodeGenerationError(f"无效的台区编码: {station_code}")

			# 生成池键：box:{station_code}
			pool_key = f"box:{station_code}"

			# 获取流水号
			sequence = await self._get_next_sequence(pool_key, self.BOX_POOL_SIZE, db)

			# 组装编码: B + 台区码(去掉T) + 4位流水
			box_code = f"{self.BOX_PREFIX}{station_code[1:]}{sequence:04d}"

			logger.info(f"生成电箱编码: {box_code}, station_code={station_code}, sequence={sequence}")
			return box_code

		except Exception as e:
			if isinstance(e, CodeGenerationError):
				raise
			logger.exception(f"生成电箱编码失败: station_code={station_code}")
			raise CodeGenerationError(f"生成电箱编码失败: {e}")

	# ============ RTU编码生成 ============

	async def generate_rtu_code(self, parent_code: str, parent_type: str = "box", db: DatabaseService = None) -> str:
		"""
		生成RTU编码（优化后）
		- 电箱下RTU: R + 9位区县码 + 3位台区流水 + 3位RTU流水
		- 台区直挂RTU: R + 9位区县码 + 台区流水(3位) + RTU流水(3位)

		Args:
			parent_code: 父级设备编码（电箱码或台区码）
			parent_type: 父级类型 ("box" 或 "station")
			db: 数据库服务实例，用于数据库降级

		Returns:
			RTU编码

		Raises:
			CodeGenerationError: 生成失败
		"""
		try:
			if parent_type == "box":
				# 电箱下RTU: R + 9位区县码 + 3位台区流水 + 3位RTU流水
				if not parent_code or len(parent_code) != self.BOX_CODE_LEN or not parent_code.startswith(
						self.BOX_PREFIX):
					raise CodeGenerationError(f"无效的电箱编码: {parent_code}")

				# 从电箱码中提取区县码和台区流水
				# 电箱码: B110105000000001 = B + 110105000(9位区县码) + 00001(5位台区流水)
				area_code = parent_code[1:10]  # 9位区县码
				station_seq = parent_code[10:15]  # 5位台区流水，取前3位
				station_seq_3d = station_seq[:3]

				pool_key = f"rtu:{parent_code}"
				sequence = await self._get_next_sequence(pool_key, 999, db)  # 3位RTU流水
				rtu_code = f"{self.RTU_PREFIX}{area_code}{station_seq_3d}{sequence:03d}"

			elif parent_type == "station":
				# 台区直挂RTU: R + 9位区县码 + 3位台区流水 + 3位RTU流水
				if not parent_code or len(parent_code) != self.STATION_CODE_LEN or not parent_code.startswith(
						self.STATION_PREFIX):
					raise CodeGenerationError(f"无效的台区编码: {parent_code}")

				# 从台区码中提取区县码和台区流水
				area_code = parent_code[1:10]  # 9位区县码
				station_seq = parent_code[10:15]  # 5位台区流水，取前3位
				station_seq_3d = station_seq[:3]

				pool_key = f"rtu:{parent_code}"
				sequence = await self._get_next_sequence(pool_key, 999, db)  # 3位RTU流水
				rtu_code = f"{self.RTU_PREFIX}{area_code}{station_seq_3d}{sequence:03d}"

			else:
				raise CodeGenerationError(f"无效的父级类型: {parent_type}")

			logger.info(
				f"生成RTU编码: {rtu_code}, parent_code={parent_code}, parent_type={parent_type}, sequence={sequence}")
			return rtu_code

		except Exception as e:
			if isinstance(e, CodeGenerationError):
				raise
			logger.exception(f"生成RTU编码失败: parent_code={parent_code}, parent_type={parent_type}")
			raise CodeGenerationError(f"生成RTU编码失败: {e}")

	# ============ 电表编码生成 ============

	async def generate_meter_code(self, rtu_code: str, db: DatabaseService = None) -> str:
		"""
		生成电表编码（优化后）: M + 9位区县码 + 3位台区流水 + 3位RTU流水 + 4位电表流水

		Args:
			rtu_code: RTU编码
			db: 数据库服务实例，用于数据库降级

		Returns:
			电表编码

		Raises:
			CodeGenerationError: 生成失败
		"""
		try:
			if not rtu_code or len(rtu_code) != self.RTU_CODE_LEN or not rtu_code.startswith(self.RTU_PREFIX):
				raise CodeGenerationError(f"无效的RTU编码: {rtu_code}")

			# 从RTU码中提取关键信息
			# RTU码: R110105000001001 = R + 110105000(9位区县码) + 001(3位台区流水) + 001(3位RTU流水)
			area_code = rtu_code[1:10]  # 9位区县码
			station_seq = rtu_code[10:13]  # 3位台区流水
			rtu_seq = rtu_code[13:16]  # 3位RTU流水

			# 生成池键：meter:{rtu_code}
			pool_key = f"meter:{rtu_code}"

			# 获取流水号
			sequence = await self._get_next_sequence(pool_key, 9999, db)  # 4位电表流水，传入db参数

			# 组装编码: M + 9位区县码 + 3位台区流水 + 3位RTU流水 + 4位电表流水 = 20位
			meter_code = f"{self.METER_PREFIX}{area_code}{station_seq}{rtu_seq}{sequence:04d}"

			logger.info(f"生成电表编码: {meter_code}, rtu_code={rtu_code}, sequence={sequence}")
			return meter_code

		except Exception as e:
			if isinstance(e, CodeGenerationError):
				raise
			logger.exception(f"生成电表编码失败: rtu_code={rtu_code}")
			raise CodeGenerationError(f"生成电表编码失败: {e}")


# 全局实例
_device_code_generator = None


def get_device_code_generator() -> DeviceCodeGenerator:
	"""获取设备编码生成器单例"""
	global _device_code_generator
	if _device_code_generator is None:
		_device_code_generator = DeviceCodeGenerator()
	return _device_code_generator
